package com.earthea.chatbot.service.impl;

import cn.hutool.core.util.IdUtil;
import com.earthea.chatbot.common.entity.BaseRsp;
import com.earthea.chatbot.entity.domain.Message;
import com.earthea.chatbot.entity.response.MessageRsp;
import com.earthea.chatbot.entity.response.PublishMessageRsp;
import com.earthea.chatbot.mapper.MessageMapper;
import com.earthea.chatbot.service.MessageService;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.model.StreamingResponseHandler;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.chat.StreamingChatLanguageModel;
import dev.langchain4j.model.output.Response;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Sinks;

import java.util.List;

@Service
public class MessageServiceImpl implements MessageService {
    @Autowired
    private MessageMapper messageMapper;

    @Autowired
    private StreamingChatLanguageModel streamingChatLanguageModel;

    @Autowired
    private ChatLanguageModel chatLanguageModel;

    @Override
    public MessageRsp selectByChat(Message message) {
        List<Message> messages = messageMapper.selectByChat(message);
        MessageRsp rsp = BaseRsp.success(MessageRsp.class);
        rsp.setMessages(messages);
        return rsp;
    }

    @Override
    public Flux<String> stream(Message message) {
        message.setMessageId(IdUtil.simpleUUID());
        messageMapper.insert(message);

        // 记录ai消息
        // 4. 创建响应处理器
        Sinks.Many<String> sink = Sinks.many().unicast().onBackpressureBuffer();

        // 5. 异步调用大模型
        streamingChatLanguageModel.generate(message.getContent(), new StreamingResponseHandler<AiMessage>() {
            private final StringBuilder fullResponse = new StringBuilder();

            @Override
            public void onNext(String token) {
                sink.tryEmitNext(token);
                fullResponse.append(token);
            }

            @Override
            public void onComplete(Response<AiMessage> response) {
                message.setRole("assistant");
                message.setMessageId(IdUtil.simpleUUID());
                message.setContent(fullResponse.toString());
                messageMapper.insert(message);
                sink.tryEmitComplete();
            }

            @Override
            public void onError(Throwable error) {
                sink.tryEmitError(error);
            }
        });
        return sink.asFlux();
    }

    @Override
    public PublishMessageRsp insert(Message message) {
        message.setMessageId(IdUtil.simpleUUID());
        int insertUser = messageMapper.insert(message);

        // 记录ai消息
        String content = chatLanguageModel.generate(message.getContent());
        message.setRole("assistant");
        message.setMessageId(IdUtil.simpleUUID());
        message.setContent(content);
        int insertAi = messageMapper.insert(message);
        // 返回ai消息
        PublishMessageRsp rsp = insertUser > 0 && insertAi > 0 ? BaseRsp.success(PublishMessageRsp.class) : BaseRsp.failed(PublishMessageRsp.class);
        rsp.setMessage(content);
        return rsp;
    }
}
